package vectorstore

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"

	"github.com/cloudwego/eino/schema"
	"github.com/sirupsen/logrus"
)

// ChromaIndexer Chroma索引器实现
type ChromaIndexer struct {
	config         *ChromaIndexerConfig
	documents      map[string]*schema.Document
	collectionPath string
	initialized    bool
}

// NewChromaIndexer 创建新的Chroma索引器
func NewChromaIndexer(config *ChromaIndexerConfig) (*ChromaIndexer, error) {
	logrus.WithField("db_path", config.DBPath).WithField("collection", config.CollectionName).Info("初始化Chroma索引器")

	collectionPath := filepath.Join(config.DBPath, config.CollectionName)
	// 确保目录存在
	if err := os.MkdirAll(collectionPath, 0755); err != nil {
		return nil, fmt.Errorf("创建集合目录失败: %w", err)
	}

	return &ChromaIndexer{
		config:         config,
		documents:      make(map[string]*schema.Document),
		collectionPath: collectionPath,
	}, nil
}

// AddDocuments 添加文档到索引
func (c *ChromaIndexer) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("Chroma索引器添加文档")

	for _, doc := range docs {
		if doc.ID == "" {
			return fmt.Errorf("文档缺少ID")
		}
		if doc.Embedding == nil || len(doc.Embedding) != c.config.Dimension {
			return fmt.Errorf("文档嵌入向量维度不匹配，期望: %d, 实际: %d", c.config.Dimension, len(doc.Embedding))
		}

		c.documents[doc.ID] = doc
	}

	c.initialized = true
	return nil
}

// UpdateDocuments 更新文档
func (c *ChromaIndexer) UpdateDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("Chroma索引器更新文档")

	// 先删除旧文档
	idsToDelete := make([]string, len(docs))
	for i, doc := range docs {
		idsToDelete[i] = doc.ID
	}

	if err := c.DeleteDocuments(ctx, idsToDelete); err != nil {
		return err
	}

	// 再添加新文档
	return c.AddDocuments(ctx, docs)
}

// DeleteDocuments 删除文档
func (c *ChromaIndexer) DeleteDocuments(ctx context.Context, ids []string) error {
	logrus.WithField("count", len(ids)).Info("Chroma索引器删除文档")

	for _, id := range ids {
		delete(c.documents, id)
	}

	return nil
}

// Search 搜索相似文档
func (c *ChromaIndexer) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	logrus.WithField("topK", topK).Info("Chroma索引器搜索文档")

	if !c.initialized {
		return []*schema.Document{}, nil
	}

	if len(queryVector) != c.config.Dimension {
		return nil, fmt.Errorf("查询向量维度不匹配，期望: %d, 实际: %d", c.config.Dimension, len(queryVector))
	}

	// 简单的余弦相似度搜索
	type scoredResult struct {
		doc   *schema.Document
		score float64
	}

	var results []scoredResult

	for _, doc := range c.documents {
		score := cosineSimilarity(queryVector, doc.Embedding)
		results = append(results, scoredResult{doc: doc, score: score})
	}

	// 按分数排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].score > results[j].score
	})

	// 返回前topK个结果
	resultDocs := []*schema.Document{}
	for i := 0; i < len(results) && i < topK; i++ {
		resultDocs = append(resultDocs, results[i].doc)
	}

	return resultDocs, nil
}

// GetDocuments 根据ID获取文档
func (c *ChromaIndexer) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	logrus.WithField("count", len(ids)).Info("Chroma索引器获取文档")

	var results []*schema.Document
	for _, id := range ids {
		if doc, exists := c.documents[id]; exists {
			results = append(results, doc)
		}
	}

	return results, nil
}

// Save 保存索引
func (c *ChromaIndexer) Save(ctx context.Context) error {
	logrus.WithField("collection_path", c.collectionPath).Info("保存Chroma索引")

	// 保存文档和嵌入向量
	collectionData := map[string]interface{}{
		documents:  c.documents,
		dimension:  c.config.Dimension,
		metricType: c.config.MetricType,
	}

	collectionDataBytes, err := json.Marshal(collectionData)
	if err != nil {
		return fmt.Errorf("序列化集合数据失败: %w", err)
	}

	dataFilePath := filepath.Join(c.collectionPath, "collection.json")
	if err := os.WriteFile(dataFilePath, collectionDataBytes, 0644); err != nil {
		return fmt.Errorf("写入集合文件失败: %w", err)
	}

	return nil
}

// Load 加载索引
func (c *ChromaIndexer) Load(ctx context.Context) error {
	logrus.WithField("collection_path", c.collectionPath).Info("加载Chroma索引")

	dataFilePath := filepath.Join(c.collectionPath, "collection.json")
	if _, err := os.Stat(dataFilePath); os.IsNotExist(err) {
		logrus.Info("集合文件不存在，使用空集合")
		return nil
	}

	dataBytes, err := os.ReadFile(dataFilePath)
	if err != nil {
		return fmt.Errorf("读取集合文件失败: %w", err)
	}

	var collectionData map[string]interface{}
	if err := json.Unmarshal(dataBytes, &collectionData); err != nil {
		return fmt.Errorf("解析集合数据失败: %w", err)
	}

	// 读取文档数据
	if docsData, ok := collectionData["documents"]; ok {
		docsBytes, err := json.Marshal(docsData)
		if err != nil {
			return fmt.Errorf("序列化文档数据失败: %w", err)
		}

		if err := json.Unmarshal(docsBytes, &c.documents); err != nil {
			return fmt.Errorf("解析文档数据失败: %w", err)
		}
	}

	c.initialized = true
	return nil
}

// Close 关闭索引器
func (c *ChromaIndexer) Close() error {
	logrus.Info("关闭Chroma索引器")
	// 在这个实现中不需要特殊的关闭操作
	return nil
}
